use heapless::spsc::Queue;
use log::{error, info};
use rayon::prelude::*;
use std::collections::HashMap;
use std::panic::{self, AssertUnwindSafe};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration; // 用于捕获 panic

// 定义消息类型（用户可以自定义）
#[derive(Debug, Clone)]
pub enum Message {
    Start,
    Data(MessageItem),
    DownLoad(ImageMessageItem),
    Stop,
}

#[derive(Debug, Clone)]
pub struct MessageItem {
    pub name: String,
    pub value: HashMap<String, String>,
}

#[derive(Debug, Clone)]
pub struct ImageMessageItem {
    pub url: String,
    pub path: String,
}

// 定义服务端配置
pub struct ServerConfig {
    pub num_threads: usize,                                // 线程数
    pub handlers: Vec<Box<dyn Fn(Message) + Send + Sync>>, // 消息处理函数列表
}

impl ServerConfig {
    pub fn new(num_threads: usize) -> Self {
        Self {
            num_threads,
            handlers: Vec::new(),
        }
    }

    // 添加消息处理函数
    pub fn add_handler<F>(&mut self, handler: F)
    where
        F: Fn(Message) + Send + Sync + 'static,
    {
        self.handlers.push(Box::new(handler));
    }
}

// 定义服务端
pub struct Server {
    queue: Arc<Mutex<Queue<Message, 1024>>>,
    config: ServerConfig,
    shutdown_flag: Arc<Mutex<bool>>, // 关闭标志位
}

impl Server {
    pub fn new(queue: Arc<Mutex<Queue<Message, 1024>>>, config: ServerConfig) -> Self {
        Self {
            queue,
            config,
            shutdown_flag: Arc::new(Mutex::new(false)), // 初始化关闭标志位
        }
    }

    pub fn start(&self) {
        info!(
            "Server started with {} threads, waiting for messages...",
            self.config.num_threads
        );

        // 克隆关闭标志位
        let shutdown_flag = Arc::clone(&self.shutdown_flag);

        // 使用 rayon 的独立线程池
        let pool = match rayon::ThreadPoolBuilder::new()
            .num_threads(self.config.num_threads)
            .build()
        {
            Ok(pool) => pool,
            Err(err) => {
                error!("Failed to create thread pool: {}", err);
                return;
            }
        };

        pool.install(|| {
            (0..self.config.num_threads).into_par_iter().for_each(|_| {
                loop {
                    // 检查关闭标志位
                    if *shutdown_flag.lock().unwrap() {
                        info!("Server: Shutdown flag detected, stopping...");
                        break;
                    }

                    // 从队列中读取消息
                    let message = {
                        let mut queue = match self.queue.lock() {
                            Ok(queue) => queue,
                            Err(_) => break,
                        };
                        queue.dequeue()
                    };

                    if let Some(message) = message {
                        // 如果收到 Stop 消息，设置关闭标志位
                        if let Message::Stop = message {
                            *shutdown_flag.lock().unwrap() = true;
                            info!("Server: Stop message received, shutting down...");
                            break;
                        }

                        // 调用所有注册的处理函数，捕获可能的 panic
                        for handler in &self.config.handlers {
                            let result = panic::catch_unwind(AssertUnwindSafe(|| {
                                handler(message.clone());
                            }));

                            if let Err(err) = result {
                                error!("Handler panicked: {:?}", err);
                            }
                        }
                    } else {
                        // 如果队列为空，休眠一段时间
                        thread::sleep(Duration::from_millis(100));
                    }
                }
            });
        });

        info!("Server stopped.");
    }

    // 关闭服务端的方法
    pub fn shutdown(&self) {
        info!("Server: Shutting down...");
        *self.shutdown_flag.lock().unwrap() = true;
    }
}

// 定义客户端
pub struct Client {
    queue: Arc<Mutex<Queue<Message, 1024>>>,
}

impl Client {
    pub fn new(queue: Arc<Mutex<Queue<Message, 1024>>>) -> Self {
        Self { queue }
    }

    // 发送消息的方法（带重试机制）
    pub fn send(&self, message: Message) {
        let mut retries = 5; // 最大重试次数
        loop {
            let mut queue = match self.queue.lock() {
                Ok(queue) => queue,
                Err(_) => break,
            };
            if queue.enqueue(message.clone()).is_ok() {
                info!("Client: Message sent successfully.");
                break;
            } else {
                retries -= 1;
                if retries == 0 {
                    error!("Client: Failed to send message after retries, queue is full!");
                    break;
                }
                info!("Client: Queue is full, retrying...");
                thread::sleep(Duration::from_millis(100)); // 等待一段时间后重试
            }
        }
    }
}
